import numpy as np
import torch
import matplotlib.pyplot as plt
from scipy.linalg import sqrtm

# GPU or CPU
device = torch.device(
    "cuda") if torch.cuda.is_available() else torch.device("cpu")

def generate_samples(n_sample, condition, guide_w):
    with torch.no_grad():
        # Load trained DDPM model
        ddpm = torch.load("DDPM/SineTest/ddpm.pth", map_location=device)
        ddpm.eval()
        x_gen, x_gen_store = ddpm.sample(n_sample=n_sample, 
                                         x_size=128, 
                                         device=device, 
                                         guide_w=guide_w, 
                                         condition=condition)
        return x_gen
    
# 定义函数计算 FID
def calculate_fid(features_real, features_fake):
    # 计算均值和协方差
    mu_real = np.mean(features_real, axis=0)
    mu_fake = np.mean(features_fake, axis=0)
    sigma_real = np.cov(features_real, rowvar=False)
    sigma_fake = np.cov(features_fake, rowvar=False)
    
    # 计算 FID
    diff = mu_real - mu_fake
    covmean = sqrtm(sigma_real @ sigma_fake)
    if np.iscomplexobj(covmean):
        covmean = covmean.real
    fid = diff @ diff + np.trace(sigma_real + sigma_fake - 2 * covmean)
    return fid

if __name__ == "__main__":
    condition=torch.tensor([1.5, 4.0, np.pi/2])#给定振幅、频率、相位
    # 生成样本
    out = generate_samples(n_sample = 30,
                         condition = condition,
                         guide_w = 0.0)
    out = out.cpu().numpy()

    print()
    
    # 创建一个图形
    plt.figure(figsize=(10, 6))

    # 遍历每一行数据并绘制曲线
    for i in range(out.shape[0]):
        plt.plot(out[i], label=f'Curve {i+1}')

    # 添加图例
    plt.legend()

    # 添加标题和轴标签
    plt.title('30 Curves with 128 Data Points Each')
    plt.xlabel('Data Points')
    plt.ylabel('Values')

    # 显示图形
    plt.show()
    